import torch
import torch.nn as nn
import torch.nn.functional as F

class Attention(torch.nn.Module):
    def __init__(self, in_dim,proj_drop_ratio):
        super(Attention, self).__init__()
        self.k = nn.Linear(in_dim, in_dim)
        self.q = nn.Linear(in_dim, in_dim)
        self.v = nn.Linear(in_dim, in_dim)
        self.proj = nn.Linear(in_dim, in_dim)
        self.proj_drop = nn.Dropout(proj_drop_ratio)
        self.GELU = nn.GELU()

    def forward(self, x):
        # k q v 均将x从in_dim转变为out_dim，特征拓展、特征对应一个权重
        k = self.k(x)
        q = self.q(x)
        v = self.v(x)
        # 点乘计算注意力
        atten = F.softmax((k * q) / torch.sqrt(torch.tensor(v.shape[1])), dim=1)
        # 特征值
        x = self.GELU(atten * v)
        x = self.proj(x)
        out = self.proj_drop(x)
        return out

class Mlp(nn.Module):
    """
    MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Block(nn.Module):
    def __init__(self,dim, mlp_ratio=4., drop_ratio=0., act_layer=nn.GELU):
        super(Block, self).__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = Attention(dim, proj_drop_ratio=drop_ratio)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.norm2 = nn.LayerNorm(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x



class Transformer(nn.Module):
    def __init__(self, num_classes, dim=8, depth=12, mlp_ratio=4.0,  drop_ratio=0.,norm_layer=None,
                 act_layer=None):
        super(Transformer, self).__init__()
        self.norm_layer = nn.LayerNorm(dim)
        act_layer = act_layer or nn.GELU
        self.blocks = nn.Sequential(*[
            Block(dim=dim, mlp_ratio=mlp_ratio,drop_ratio=drop_ratio, act_layer=act_layer)
            for i in range(depth)
        ])
        self.fc_class = nn.Linear(dim,out_features=num_classes)
        # self.sigmoid = nn.Sigmoid()
        self.apply(_init_vit_weights)

    def forward(self, x):
        x = self.blocks(x)
        x = self.norm_layer(x)
        x = self.fc_class(x)
        return x


def _init_vit_weights(m):
    """
    ViT weight initialization
    :param m: module
    """
    if isinstance(m, nn.Linear):
        nn.init.trunc_normal_(m.weight, std=.01)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode="fan_out")
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
        nn.init.zeros_(m.bias)
        nn.init.ones_(m.weight)


def model(in_dim, num_classes: int = 100):
    model = Transformer(dim=in_dim,depth=12,num_classes=num_classes)
    return model

if __name__ == '__main__':
    x = torch.randn(3,8)
    print(x)
    model = model(8,2)
    print(model(x))



